# -*- coding: utf-8 -*- #
# -----------------------------------------------------------------------
# File Name:    test.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了模型的测试流程
# -----------------------------------------------------------------------
import torch
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from dataset import CustomDataset

def test(dataloader, model, device):
    """定义测试流程（遵循指导书6-40至6-45段要求）"""
    model.eval()
    size = len(dataloader.dataset)
    correct_num = 0
    with torch.no_grad():
        for data in dataloader:
            inputs, labels = data['image'].to(device), data['label'].to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            correct_num += (predicted == labels).sum().item()
    accuracy = 100 * correct_num / size
    print(f'Test Accuracy: {accuracy}%')

if __name__ == "__main__":
    # 直接创建数据加载器和模型（不依赖pytest夹具）
    model = torch.load('./models/model.pkl', weights_only=False)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    test_dataset = CustomDataset('./images/test.txt', './images/test', ToTensor)
    test_dataloader = DataLoader(test_dataset, batch_size=32)
    # 调用test函数并传递参数
    test(test_dataloader, model, device)